from torch import nn

# 定义多层感知机模型
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.flatten = nn.Flatten()  # 定义数据展平层,将二维图像数据转换为一维向量
        self.fc1 = nn.Linear(784, 20)  # 定义第一个全连接层,输入维度为784,输出维度为20
        self.relu = nn.ReLU()  # 定义ReLU激活函数
        self.fc2 = nn.Linear(20, 10)  # 定义第二个全连接层,输入维度为20,输出维度为10

    def forward(self, x):
        x = self.flatten(x)  # 将输入数据展平为一维向量
        x = self.fc1(x)  # 通过第一个全连接层
        x = self.relu(x)  # 应用ReLU激活函数
        x = self.fc2(x)  # 通过第二个全连接层并输出结果
        return x
